import argparse
import os
import sys
import re
import json
from pathlib import Path

from c_header_extract import extract_definitions_from_header
from c_symbol_table import construct_symbol_table
from extract_rust_interface import extract_dependent_rust_interface
from rebuilt_header import rebuilt_header_information
from extract_function_symbols import extract_symbols_entry

from utils import system_header_paths, collect_include_paths

def collect_files(src_dir):
    h_files = dict()
    c_files = dict()
    for root, dirs, files in os.walk(src_dir):
        for _file in files:
            filename = os.path.join(root, _file)
            if filename.endswith('.h'):
                fname = filename[:-2]
                fname = fname.split(os.path.sep)[-1]
                absolute_file_path = filename
                content = ""
                with open(absolute_file_path, "r") as hfile:
                    content = hfile.read()
                
                h_files[fname] = (absolute_file_path, content)
            elif filename.endswith('.c'):
                fname = filename[:-2]
                fname = fname.split(os.path.sep)[-1]
                absolute_file_path = filename
                content = ""
                with open(absolute_file_path, "r") as cfile:
                    content = cfile.read()

                c_files[fname] = (absolute_file_path, content)
            elif filename.endswith(".cpp"):
                fname = filename[:-4]
                fname = fname.split(os.path.sep)[-1]
                absolute_file_path = filename
                content = ""
                with open(absolute_file_path, "r") as cppfile:
                    content = cppfile.read()

                c_files[fname] = (absolute_file_path, content)
            else:
                print(f"Unsupported file extensions{filename}!")
    return h_files, c_files

def rebuilt_header_entry(filename, header_definitions, symbol_table, output_directory, include_dirs):
    header_definitions, symbol_table = rebuilt_header_information(filename, header_definitions, symbol_table, output_directory, include_dirs)
    
    return header_definitions, symbol_table

def parse_src_file(c_file, output_directory, include_dirs):
    from parse_cfg_snippets import cfg_entry
    from parse_ddg import ddg_entry
    with open(c_file, "r") as src_file:
        src_content = src_file.read()
        needed_restore = False
    try:
        cfg_entry(c_file, output_directory, include_dirs)
    except Exception as e:
        print(f"Error occured when extracting cfg of {c_file}:{e}")
        needed_restore = True
    finally:
        if needed_restore:
            with open(c_file, "w") as src_file:
                src_file.write(src_content)

    try:
        ddg_entry(c_file, output_directory, include_dirs)
    except Exception as e:
        print(f"Error occured when extracting ddg of {c_file}:{e}")

def write_header_information(filename, symbol_table, definitions, output_dir, global_defs=False):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if isinstance(definitions, list) and not global_defs:
        #header_definitions = "\n".join(definitions)
        header_definitions = definitions
    else:
        header_definitions = definitions

    output_path = Path(output_dir) / (filename + "_output.json")
    data = {
        "C": {
            "symbol_table": symbol_table,
            "header_definitions": header_definitions,
        }, 
        "Rust": {
            "symbol_table": dict(),
        },
    }

    with open(output_path, "w") as f:
        json.dump(data, f, indent=4)

def write_source_information(filename, symbol_table, header_definitions, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    output_path = Path(output_dir) / (filename + "_output.json")
    if os.path.exists(output_path):
        with open(output_path, "r") as jfile:
            data = json.load(jfile)
            if "C" in data:
                data["C"]["symbol_table"] = symbol_table
                data["C"]["header_definitions"] = header_definitions
            else:
                data["C"] = dict()
                data["C"]["symbol_table"] = symbol_table
                data["C"]["header_definitions"] = header_definitions
            data["Rust"] = dict()
            data["Rust"]["symbol_table"] = dict()
    else:
        data = {
            "C": {
                "symbol_table": symbol_table,
                "header_definitions": header_definitions,
            },
            "Rust": {
                "symbol_table": dict(),
            }
        }

    with open(output_path, "w") as f:
        json.dump(data, f, indent=4)

def write_symbols_information(filename, func_symbols, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    output_path = Path(output_dir) / (filename + "_output.json")
    if os.path.exists(output_path):
        with open(output_path, "r") as jfile:
            data = json.load(jfile)
            if "C" in data:
                if "code" in data["C"]:
                    code_snippet = data["C"]["code"]
                    idx_lst = sorted([int(idx) for idx in code_snippet.keys()])
                    for item in func_symbols:
                        if len(item) == 2:
                            func_name = item[0]
                            fsymbols = item[1]
                        else:
                            continue
                        for idx in idx_lst:
                            func_info = code_snippet[str(idx)]
                            if func_info["kind"] != 8:
                                continue
                            if func_name == func_info["name"]:
                                func_info['symbols'] = fsymbols
                else:
                    data["C"]["code"] = dict()
            else:
                data["C"] = dict()
    else:
        data = {
            "C": {
            },
            "Rust": {
                "symbol_table": dict(),
            }
        }

    with open(output_path, "w") as f:
        json.dump(data, f, indent=4)


def parse_header_file(header_file, include_dirs):
    symbol_table = construct_symbol_table(header_file, include_dirs)

    return symbol_table

def extract_rust_interface_entry(filename, header_file, output_dir, include_dirs):
    rust_interface, pro_include_files = extract_dependent_rust_interface(header_file, include_dirs)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    output_path = Path(output_dir) / (filename + "_output.json")
    
    if len(rust_interface) > 0:
        if os.path.exists(output_path):
            with open(output_path, "r") as jfile:
                data = json.load(jfile)
                if "Rust" in data and "symbol_table" in data["Rust"]:
                    data["Rust"]["symbol_table"] = rust_interface
                else:
                    # we add the Rust table.
                    data["Rust"] = dict()
                    data["Rust"]["symbol_table"] = rust_interface
            with open(output_path, "w") as f:
                json.dump(data, f, indent=4)
    if len(pro_include_files) > 0:
        if os.path.exists(output_path):
            with open(output_path, "r") as jfile:
                data = json.load(jfile)
                data["include_files"] = pro_include_files
        else:
            data = dict()
            data["include_files"] = pro_include_files
        with open(output_path, "w") as f:
            json.dump(data, f, indent=4)
    if "framework" in filename:
        if os.path.exists(output_path):
            with open(output_path, "r") as jfile:
                data = json.load(jfile)
                data["include_files"] = ["alloc-testing_output.json"]
        else:
            data = dict()
            data["include_files"] = ["alloc-testing_output.json"] 
        with open(output_path, "w") as f:
            json.dump(data, f, indent=4)


def preprocess_entry(c_dir, c_src_dir, output_dir):
    # 使用argparse解析命令行参数
    #parser = argparse.ArgumentParser(description="Extract C file pairs")
    #parser.add_argument("c_dir", type=str, help="Pass a C directory path to be parsed.")
    #parser.add_argument("c_src_dir", type=str, help="Pass a C src path to be parsed.")
    #parser.add_argument("output_directory", type=str, help="Pass a C include path to be parsed.")
    #args = parser.parse_args()

    #output_dir = args.output_directory
    #c_dir = args.c_dir
    #c_src_dir = os.path.join(c_dir, args.c_src_dir)
    c_src_dir = os.path.join(c_dir, c_src_dir) 
    # 检查C directory 是否存在
    if not os.path.isdir(c_dir):
        print(f"Error: code directory {c_dir} dose not exist!")
        return
    elif not os.path.isdir(c_src_dir):
        print(f"Error: src directory {c_src_dir} does not exist!")
        return
    
    include_dirs = collect_include_paths(c_dir)
    print(f"Scan header path in {c_src_dir}")
    if include_dirs:
        print("Found header path:")
        print('\n'.join(include_dirs))
    else:
        print(f"Use system header path!")

    h_files, c_files = collect_files(c_src_dir)
    include_dirs.extend(system_header_paths)
    # parse header and src files.
    for h_filename in h_files.keys():
        if h_filename in c_files:
            print(f"Parsing {h_filename}.h and {h_filename}.c files!")
            symbol_table = parse_header_file(c_files[h_filename][0], include_dirs)
            header_definitions = extract_definitions_from_header(c_files[h_filename][0], include_dirs)
            write_header_information(h_filename, symbol_table, header_definitions, output_dir)
            parse_src_file(c_files[h_filename][0], output_dir, include_dirs)
            header_definitions, symbol_table = rebuilt_header_entry(h_filename, header_definitions, symbol_table, output_dir, include_dirs)
            func_symbols = extract_symbols_entry(c_files[h_filename][0], include_dirs)
            write_symbols_information(h_filename, func_symbols, output_dir)
            if c_src_dir == "test":
                extract_rust_interface_entry(h_filename, c_files[h_filename][0], output_dir, include_dirs)
            c_files.pop(h_filename)
        else:
            # parse single .h files
            print(f"Parsing {h_filename}.h file!")
            symbol_table = parse_header_file(h_files[h_filename][0], include_dirs)
            header_definitions = extract_definitions_from_header(h_files[h_filename][0], include_dirs)
            write_header_information(h_filename, symbol_table, header_definitions, output_dir)
    # parse single .c/.cpp files.
    for c_filename in c_files.keys():
        print(f"Parsing single {c_filename}.c file!")
        parse_src_file(c_files[c_filename][0], output_dir, include_dirs)
        src_symbol_table = parse_header_file(c_files[c_filename][0], include_dirs)
        src_header_definitions = extract_definitions_from_header(c_files[c_filename][0], include_dirs)
        write_source_information(c_filename, src_symbol_table, src_header_definitions, output_dir)
        header_definitions, symbol_table = rebuilt_header_entry(c_filename, src_header_definitions, src_symbol_table, output_dir, include_dirs)
        func_symbols = extract_symbols_entry(c_files[c_filename][0], include_dirs)
        write_symbols_information(h_filename, func_symbols, output_dir)
        extract_rust_interface_entry(c_filename, c_files[c_filename][0], output_dir, include_dirs)

def test_c_files_preprocess():
    import sys
    # 模拟命令行参数
    test_c_dir = "test_files"
    test_c_src_dir = "src"
    test_output_dir = "test_output"
    
    #sys.argv = [
    #    "c_files_preprocess.py",
    #    test_c_dir,
    #    test_c_src_dir,
    #    test_output_dir
    #]
    
    preprocess_entry(test_c_dir, test_c_src_dir, test_output_dir)

if __name__ == "__main__":
    test_c_files_preprocess()
